package net.aias.springaichat.adapter.ai.client;

import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.ArrayList;
import java.util.List;

/**
 * @Auther: cheng.tang
 * @Date: 2025/5/13
 * @Description: spring-ai-chat
 */

public class DoubaoEmbeddingClient {

    private final WebClient webClient;

    public DoubaoEmbeddingClient(WebClient webClient) {
        this.webClient = webClient;
    }

    public EmbeddingResponse embed(List<String> texts) {
        EmbeddingRequest request = new EmbeddingRequest();
        request.setModel("doubao-embedding-text-240715"); // 使用示例中的模型
        request.setInput(texts);
        request.setEncodingFormat("float"); // 添加编码格式

        DoubaoApiResponse response = webClient.post()
                .bodyValue(request)
                .retrieve()
                .bodyToMono(DoubaoApiResponse.class)
                .block();

        List<Embedding> embeddings = new ArrayList<>();
        if (response != null && !response.getData().isEmpty()) {
            for (DoubaoApiResponse.Data data : response.getData()) {
                embeddings.add(new Embedding(data.getEmbedding(), data.getIndex()));
            }
        }
        return new EmbeddingResponse(embeddings);
    }

    // 请求模型（包含encoding_format）
    private static class EmbeddingRequest {
        private String model;
        private List<String> input;
        private String encoding_format = "float"; // 默认使用float格式

        public String getModel() {
            return model;
        }

        public void setModel(String model) {
            this.model = model;
        }

        public List<String> getInput() {
            return input;
        }

        public void setInput(List<String> input) {
            this.input = input;
        }

        public String getEncodingFormat() {
            return encoding_format;
        }

        public void setEncodingFormat(String encodingFormat) {
            this.encoding_format = encodingFormat;
        }
    }

    // 响应模型（简化版，仅保留必要字段）
    private static class DoubaoApiResponse {
        private List<Data> data;

        public List<Data> getData() {
            return data;
        }

        public static class Data {
            private float[] embedding;
            private int index;

            public float[] getEmbedding() {
                return embedding;
            }

            public int getIndex() {
                return index;
            }
        }
    }


}
